
import sys
import numpy as np
import pandas as pd
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                               QHBoxLayout, QPushButton, QComboBox, QLabel,
                               QSpinBox, QGroupBox, QMessageBox)
from PySide6.QtCore import Qt, QDateTime
import pyqtgraph as pg
from pyqtgraph import DateAxisItem
import psycopg2
from sqlalchemy import create_engine

# 设置PyQtGraph使用Qt6
pg.setConfigOption('useOpenGL', True)
pg.setConfigOption('antialias', True)


class CandlestickItem(pg.GraphicsObject):
    def __init__(self):
        pg.GraphicsObject.__init__(self)
        self.data = None
        self.generatePicture()

    def set_data(self, data):
        self.data = data
        self.generatePicture()
        self.update()

    def generatePicture(self):
        self.picture = pg.QtGui.QPicture()
        p = pg.QtGui.QPainter(self.picture)

        if self.data is None:
            p.end()
            return

        w = (self.data[1][0] - self.data[0][0]) / 3.0
        for (t, open, close, min, max) in self.data:
            # 绘制K线
            pen = pg.mkPen('g' if close > open else 'r', width=2)
            p.setPen(pen)

            # 绘制影线
            p.drawLine(pg.QtCore.QPointF(t, min), pg.QtCore.QPointF(t, max))

            # 绘制实体
            if close > open:
                brush = pg.mkBrush('g')
            else:
                brush = pg.mkBrush('r')
            p.setBrush(brush)

            p.drawRect(pg.QtCore.QRectF(t - w, open, w * 2, close - open))

        p.end()

    def paint(self, p, *args):
        p.drawPicture(0, 0, self.picture)

    def boundingRect(self):
        return pg.QtCore.QRectF(self.picture.boundingRect())


class StockKLineChart(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("股票K线图")
        self.setGeometry(100, 100, 1200, 800)
        self.select_tick=f'''select "InstrumentID", "DateTime", coalesce("PreSettlementPrice",0.), coalesce("PreClosePrice",0.), coalesce("PreOpenInterest",0.), 
	"UpperLimitPrice", "LowerLimitPrice","OpenPrice", 	"HighestPrice", "LowestPrice", 
	 "Volume", "Turnover", "AveragePrice" , "OpenInterest", "LastPrice",
	"BidPrice1", "BidVolume1", "AskPrice1", "AskVolume1"
	from tick.tmpl where "InstrumentID"=$1 and "DateTime" between $2  and $3 order by "DateTime" desc limit 10;
	'''
        self.select_bar=f'''SELECT date_time, instrument_id, trading_day, open, high, low, close, volume, open_interest, pre_vol, ticks, ins_dt, exchange
	FROM bar.bar
	order by date_time desc
	limit 1000;'''

        # 数据库连接参数
        self.db_params = {
            'host': '192.168.10.5',
            'database': 'ctp',
            'user': 'postgres',
            'password': 'steward1!',
            'port': '5432'
        }

        # 创建主窗口部件
        self.main_widget = QWidget()
        self.setCentralWidget(self.main_widget)

        # 创建主布局
        self.main_layout = QVBoxLayout(self.main_widget)

        # 创建控制面板
        self.create_control_panel()

        # 创建图表
        self.create_chart()

        # 初始化数据
        self.stock_data = None
        self.load_stock_data()

    def create_control_panel(self):
        # 控制面板组
        control_group = QGroupBox("控制面板")
        control_layout = QHBoxLayout()

        # 股票选择
        stock_label = QLabel("股票代码:")
        self.stock_combo = QComboBox()
        self.stock_combo.addItems(["000001.SZ", "000002.SZ", "600001.SH"])
        self.stock_combo.currentTextChanged.connect(self.load_stock_data)

        # 时间周期选择
        period_label = QLabel("时间周期:")
        self.period_combo = QComboBox()
        self.period_combo.addItems(["日线", "周线", "月线"])
        self.period_combo.currentTextChanged.connect(self.load_stock_data)

        # 显示天数
        days_label = QLabel("显示天数:")
        self.days_spin = QSpinBox()
        self.days_spin.setRange(30, 1000)
        self.days_spin.setValue(250)
        self.days_spin.setSingleStep(10)

        # 更新按钮
        self.update_button = QPushButton("更新")
        self.update_button.clicked.connect(self.update_chart)

        # 数据库连接状态
        self.db_status_label = QLabel("数据库状态: 未连接")

        # 添加控件到布局
        control_layout.addWidget(stock_label)
        control_layout.addWidget(self.stock_combo)
        control_layout.addWidget(period_label)
        control_layout.addWidget(self.period_combo)
        control_layout.addWidget(days_label)
        control_layout.addWidget(self.days_spin)
        control_layout.addWidget(self.update_button)
        control_layout.addWidget(self.db_status_label)
        control_layout.addStretch()

        control_group.setLayout(control_layout)
        self.main_layout.addWidget(control_group)

    def create_chart(self):
        # 创建图表部件
        self.chart_widget = pg.PlotWidget(axisItems={'bottom': DateAxisItem()})
        self.chart_widget.setBackground('w')
        self.chart_widget.showGrid(x=True, y=True, alpha=0.3)

        # 创建K线图项
        self.candlestick_item = CandlestickItem()
        self.chart_widget.addItem(self.candlestick_item)

        # 创建成交量图
        self.volume_widget = pg.PlotWidget(axisItems={'bottom': DateAxisItem()})
        self.volume_widget.setMaximumHeight(200)
        self.volume_widget.setBackground('w')
        self.volume_widget.showGrid(x=True, y=True, alpha=0.3)

        # 创建成交量柱状图
        self.volume_item = pg.BarGraphItem(x=[], height=[], width=0.6, brush='r')
        self.volume_widget.addItem(self.volume_item)

        # 添加图表到主布局
        self.main_layout.addWidget(self.chart_widget)
        self.main_layout.addWidget(self.volume_widget)

        # 设置十字光标
        self.vLine = pg.InfiniteLine(angle=90, movable=False)
        self.hLine = pg.InfiniteLine(angle=0, movable=False)
        self.chart_widget.addItem(self.vLine, ignoreBounds=True)
        self.chart_widget.addItem(self.hLine, ignoreBounds=True)

        self.proxy = pg.SignalProxy(self.chart_widget.scene().sigMouseMoved,
                                    rateLimit=60, slot=self.mouse_moved)

    def get_db_connection(self):
        """创建数据库连接"""
        try:
            # 使用SQLAlchemy创建引擎
            engine = create_engine(
                f"postgresql+psycopg2://{self.db_params['user']}:{self.db_params['password']}"
                f"@{self.db_params['host']}:{self.db_params['port']}/{self.db_params['database']}"
            )
            return engine
        except Exception as e:
            QMessageBox.critical(self, "数据库错误", f"连接数据库失败: {str(e)}")
            return None

    def load_stock_data(self):
        """从数据库加载股票数据"""
        engine = self.get_db_connection()
        if engine is None:
            return

        try:
            # 获取当前选择的股票代码和时间周期
            instrument_id = self.stock_combo.currentText()
            period = self.period_combo.currentText()

            # 根据时间周期确定表名
            table_map = {
                "日线": "bar",
                "周线": "weekly_data",
                "月线": "monthly_data"
            }
            table_name = table_map.get(period, "daily_data")
            instrument_id='v2509'

            # 构建SQL查询
            query = f"""
                SELECT date_time as date, open, close, low, high, volume 
                FROM bar.{table_name} 
                WHERE instrument_id = '{instrument_id}'
                ORDER BY date_time ASC
                limit 100;
            """
            print(query)

            # 使用pandas读取数据
            df = pd.read_sql(query, engine)

            if df.empty:
                QMessageBox.warning(self, "数据警告", "未找到相关股票数据")
                return

            # 将日期转换为时间戳
            df['date'] = pd.to_datetime(df['date']).astype(np.int64) // 10 ** 9

            # 转换为numpy数组
            self.stock_data = df.to_records(index=False)

            # 更新数据库状态
            self.db_status_label.setText("数据库状态: 已连接")
            self.db_status_label.setStyleSheet("color: green")

            # 更新图表
            self.update_chart()

        except Exception as e:
            QMessageBox.critical(self, "数据加载错误", f"加载数据失败: {str(e)}")
            self.db_status_label.setText("数据库状态: 连接错误")
            self.db_status_label.setStyleSheet("color: red")

    def update_chart(self):
        if self.stock_data is None:
            return

        # 获取显示天数
        days = self.days_spin.value()

        # 获取最近的数据
        data = self.stock_data[-days:]
        print(data[-10:])

        # 更新K线图
        kline_data = np.column_stack([
            data['date'],
            data['open'],
            data['close'],
            data['low'],
            data['high']
        ])
        self.candlestick_item.set_data(kline_data)

        # 更新成交量图
        volume_colors = ['g' if close > open else 'r'
                         for close, open in zip(data['close'], data['open'])]
        self.volume_item.setOpts(
            x=data['date'],
            height=data['volume'] / 1000000,  # 转换为百万单位
            brushes=volume_colors
        )

        # 更新坐标轴范围
        self.chart_widget.setXRange(data['date'][0], data['date'][-1])
        self.volume_widget.setXRange(data['date'][0], data['date'][-1])

        # 自动调整Y轴范围
        self.chart_widget.enableAutoRange(axis='y')
        self.volume_widget.enableAutoRange(axis='y')

    def mouse_moved(self, evt):
        pos = evt[0]  # 获取鼠标位置
        if self.chart_widget.sceneBoundingRect().contains(pos):
            mouse_point = self.chart_widget.plotItem.vb.mapSceneToView(pos)
            self.vLine.setPos(mouse_point.x())
            self.hLine.setPos(mouse_point.y())


def main():
    app = QApplication(sys.argv)
    window = StockKLineChart()
    window.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main()
'''
这个更新后的版本添加了以下功能：

1. ** 数据库连接 **：
- 使用SQLAlchemy创建PostgreSQL连接
- 支持配置数据库连接参数（主机、数据库名、用户名、密码、端口）

2. ** 数据读取 **：
- 使用pandas.read_sql()
从PostgreSQL读取数据
- 支持不同时间周期的数据表（日线、周线、月线）
- 自动将日期转换为时间戳格式

3. ** 错误处理 **：
- 数据库连接错误处理
- 数据加载错误处理
- 空数据检查

4. ** 状态显示 **：
- 显示数据库连接状态
- 连接成功显示绿色，错误显示红色

要使用这个程序，你需要：

1.
安装必要的依赖：
```bash
pip
install
PySide6
pyqtgraph
numpy
pandas
psycopg2
sqlalchemy
```

2.
确保你的PostgreSQL数据库中有相应的表结构。例如，日线数据表应该包含以下字段：
```sql
CREATE
TABLE
daily_data(
    stock_code
VARCHAR(10),
date
DATE,
open
NUMERIC,
high
NUMERIC,
low
NUMERIC,
close
NUMERIC,
volume
BIGINT,
PRIMARY
KEY(stock_code, date)
);
```

3.
在程序中修改数据库连接参数（host、database、user、password、port）以匹配你的PostgreSQL配置。

4.
确保数据库中已经包含了股票数据。你可以根据需要修改SQL查询语句以适应你的数据库表结构。

这个版本提供了更实用的功能，可以直接从数据库中读取真实的股票数据进行显示。
'''